from typing import List, Optional

from bson import ObjectId
from pydantic import BaseModel
from pymongo import DESCENDING, ASCENDING


def pymongo_sort(fields):
    sort = list()
    for field in fields:
        if field.startswith("-"):
            field = field.lstrip("-")
            sort.append((field, DESCENDING))
        else:
            sort.append((field, ASCENDING))
    return sort


def shell_sort(fields):
    sort = {}
    for field in fields:
        if field.startswith("-"):
            field = field.lstrip("-")
            sort[field] = DESCENDING
        else:
            sort[field] = ASCENDING
    return sort


def convert2sort(fields, is_shell=False):
    if not fields:
        return None
    if is_shell:
        return shell_sort(fields)
    return pymongo_sort(fields)


def convert2projection(fields):
    if not fields:
        return None
    projection = dict()
    for field in fields:
        if field.startswith("-"):
            field = field.lstrip("-")
            projection[field] = False
        else:
            projection[field] = True
    return projection


def unify_doc(doc):
    oid = doc.pop('_id')
    doc['id'] = str(oid)
    return doc


def add_fields_opt_for_order(field, order_field):
    """
    返回用于排序的新字段，确保空值排在最后面
    :param field:
    :param order_field:
    :return:
    """
    return {
        order_field: {
            "$function": {
                "body": "function(value){const isEmptyObj=function(obj){for(let name in obj){return false}return true};const isEmptyArray=function(arr){return arr.length===0};if(value===null||value===undefined||value===''){return false}if(typeof value==='object'){if(Array.isArray(value)){return!isEmptyArray(value)}return!isEmptyObj(value)}return true}",
                "args": [f"${field}"],
                "lang": "js"
            }
        }
    }


def get_sort_stages_for_order(sort: list):
    first_field = sort[0]
    if first_field.startswith("-"):
        first_field = first_field.lstrip("-")
    order_field = f"{first_field}_exists"
    sort = [f"-{order_field}"] + sort
    stages = [
        {"$addFields": add_fields_opt_for_order(first_field, order_field)},
        {"$sort": convert2sort(sort, is_shell=True)}
    ]
    return stages


class LookupConfig(BaseModel):
    from_: str
    local_field: str
    foreign_field: str
    as_field: str
    fields: Optional[List[str]]
    ref_type: str  # 取值["list", "object"]


def get_lookup_stages(configs: List[LookupConfig]):
    stages = []
    add_fields_stage = {}
    for conf in configs:
        pipelines = []
        stage = {
            "from": conf.from_,
            "localField": conf.local_field,
            "foreignField": conf.foreign_field,
        }
        if conf.ref_type == 'list':
            stage.update()
            stage['as'] = conf.as_field
            stage['let'] = {"order_ids": f"${conf.local_field}"}
            # 保持list中元素的顺序
            pipelines.extend([
                {"$addFields": {"order": {"$indexOfArray": ["$$order_ids", "$_id"]}}},
                {"$sort": {"order": 1}},
            ])
        elif conf.ref_type == 'object':
            temp_field = f"{conf.as_field}_ref"
            stage['as'] = temp_field
            add_fields_stage[conf.as_field] = {"$first": f"${temp_field}"}
            add_fields_stage[temp_field] = "$$REMOVE"
        else:
            raise ValueError(f"ref_type should be 'list' or 'object', get {conf.ref_type}")
        projections = {"id": "$_id"}
        if conf.fields:
            projections.update(convert2projection(conf.fields))
        pipelines.append({"$project": projections})
        stage['pipeline'] = pipelines
        stages.append({"$lookup": stage})
    if add_fields_stage:
        stages.append({"$addFields": add_fields_stage})
    return stages


class PyObjectId(ObjectId):
    @classmethod
    def __get_validators__(cls):
        yield cls.validate

    @classmethod
    def validate(cls, v):
        try:
            if not ObjectId.is_valid(v):
                raise ValueError("Invalid objectid")
            return ObjectId(v)
        except ValueError:
            return str(v)

    @classmethod
    def __modify_schema__(cls, field_schema):
        field_schema.update(type="string")

